# scripts/download_data.py
from datasets import load_dataset
import os

# ========================
# 配置
# ========================
OUTPUT_DIR = "../data"
os.makedirs(OUTPUT_DIR, exist_ok=True)

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['HF_DATASETS_TRUST_REMOTE_CODE'] = '1'  # 允许运行远程脚本
os.environ['HF_DATASETS_CACHE'] = './hf_cache'     # 指定缓存目录，方便查看

TEXT_FILE = os.path.join(OUTPUT_DIR, "openwebtext_sample.txt")
MAX_LINES = 250_000        # 取25万条高质量文本
MIN_LEN = 100             # 过滤太短的文本

# ========================
# 下载 OpenWebText 子集
# ========================
print("🚀 正在加载 OpenWebText 数据集（流式加载）...")

# 使用 streaming=True 避免全量加载
dataset = load_dataset("openwebtext", split="train", streaming=True, trust_remote_code=True)

print(f"📝 正在写入文本到 {TEXT_FILE} ...")
with open(TEXT_FILE, "w", encoding="utf-8") as f:
    count = 0
    for example in dataset:
        text = example["text"].strip()
        if len(text) >= MIN_LEN: # 过滤长度不足 100 个字符的短文本
            # 使用一个换行符分隔不同文档（类似段落分隔）
            if count > 0:
                f.write("\n")  # 在文档之间添加空行分隔
            f.write(text)
            count += 1

            if count % 10_000 == 0:
                print(f"✅ 已保存 {count} 条")

        if count >= MAX_LINES:
            break

print(f"🎉 数据下载完成！共 {count} 条，保存至: {TEXT_FILE}")